import struct
import VoxDataset
import torch
from torch.utils.data import Dataset
import numpy as np
import os

class Visualization:
    def __init__(self, originpath, rebuildpath, num):
        self.originpath = originpath
        self.rebuildpath = rebuildpath
        self.num = num
        with open(self.rebuildpath, 'rb') as f:
            self.rebuilddata = VoxDataset.VoxDataset(f)
        with open(self.originpath, 'rb') as f:
            self.origindata = VoxDataset.VoxDataset(f)

    def outputPC(self):
        with open("模型/可视化模型/" + self.num.__str__() + ".txt", 'w') as fp:
            for i in range(self.origindata.__len__()):
                coordinate, flag = self.origindata.__getitem__(i)
                coordinate, flag2 = self.rebuilddata.__getitem__(i)
                RGB = self.getRGB(flag,flag2)
                if (RGB[0] == 1):
                    fp.write(str(int(coordinate[0].item())))
                    fp.write(" ")
                    fp.write(str(int(coordinate[1].item())))
                    fp.write(" ")
                    fp.write(str(int(coordinate[2].item())))
                    fp.write(" ")
                    fp.write(str(RGB[1]))
                    fp.write(" ")
                    fp.write(str(RGB[2]))
                    fp.write(" ")
                    fp.write(str(RGB[3]))
                    fp.write("\n")

    def getRGB(self, flag, flag2):
        # 预测正确为绿
        # 预测假到真为红
        # 预测真到假为蓝
        RGB = [0,0,0,0]
        if(flag + flag2 >= 1):
            RGB[0] = 1
            if(flag == flag2):
                RGB[2] = 255
            if(flag > flag2):
                RGB[1] = 255
            if(flag < flag2):
                RGB[3] = 255
        return RGB